#! /usr/bin/python
# -*- coding: utf-8 -*-
# @Time    : 2017/11/3 11:32
# @Author  : Deyu.Tian

import os
import glob

from mpl_toolkits.mplot3d import Axes3D#for matplotlib 2.0.x
import matplotlib.pyplot as plt, matplotlib.cm as cm,  matplotlib.patches as mpatches
from image2gif import writeGif
from PIL import Image as PIL_Image

import numpy as np

import gdal

import config
MuzDir = config.muztagData
AST14DMO_DEM = MuzDir + "/AST14DMO_00304102017055823_DEM_resamp.tif"
Gaiz_GDEM = MuzDir + "/GaizDEM_toUTM_resamp.tif"
STRM_DEM = MuzDir + "/STRM_DEM_toUTM_resamp.tif"
FigDir = config.figDir

def reprojection(geo_dem):
    """
    gdal reprojection
    :param dem:
    :return:
    """
    cmd = "gdalwarp -t_srs '+proj=utm +zone=43N +datum=WGS84' -dstnodata -9999 -overwrite {} {}_toUTM.tif".format(geo_dem, geo_dem[:-4])
    os.system(cmd)

def resampling(utm_dem):
    """
    down sampling dem data for better visualize
    :param utm_dem:
    :return:
    """
    cmd = "gdal_translate -tr 500 500 -r cubic -a_nodata 0 -stats {} {}_resamp.tif".format(utm_dem, utm_dem[:-4])
    os.system(cmd)

def dem2array(dem1, dem2):
    """

    read dems to array by gdal

    :param imgfn path of geotiff

    :return narray of geotiff

    """

    dem1_data = gdal.Open(dem1)
    dem1_array = dem1_data.ReadAsArray()

    dem2_data = gdal.Open(dem2)
    dem2_array = dem2_data.ReadAsArray()

    return dem1_array, dem2_array


def dem3array(dem1, dem2, dem3):
    """

    read dems to array by gdal

    :param imgfn path of geotiff

    :return narray of geotiff

    """

    dem1_data = gdal.Open(dem1)
    dem1_array = dem1_data.ReadAsArray()

    dem2_data = gdal.Open(dem2)
    dem2_array = dem2_data.ReadAsArray()

    dem3_data = gdal.Open(dem3)
    dem3_array = dem3_data.ReadAsArray()

    return dem1_array, dem2_array, dem3_array

def scatter3d_master():
    """
    plot input geotiff as 3d scatters
    :return:
    """

    arr1, arr2, arr3 = dem3array(STRM_DEM, Gaiz_GDEM, AST14DMO_DEM)
    arr1[arr1 <= 0] = 2800 #eliminate -9999 affect visulize
    arr2[arr2 <= 0] = 2800
    arr3[arr3 <= 0] = 2800


    print(arr1.shape, np.min(arr1))#(1095, 836)
    print(arr2.shape, np.min(arr2))#(1179, 899)

    #row_slice = np.arange(0, arr1.shape[0], 2)
    #arr1_samp = np.delete(arr1, row_slice, axis=0)
    #col_slice = np.arange(0, arr1.shape[1], 2)
    #arr1_samp = np.delete(arr1_samp, col_slice, axis=1)
    #print(arr1_samp.shape)
    #return

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    zs = arr1.flatten()
    xs = np.zeros_like(zs)
    xs_copy = np.zeros_like(xs)
    ys = np.zeros_like(zs)

    for i in range(arr1.shape[0]):
        _tmp = np.full_like(xs[arr1.shape[1] * i:arr1.shape[1] * (i + 1)], i)
        ys[arr1.shape[1] * i:arr1.shape[1] * (i + 1)] = np.arange(0, arr1.shape[1], 1)
        xs_copy[arr1.shape[1] * i:arr1.shape[1] * (i + 1)] = _tmp

    ax.scatter(xs_copy, ys, zs, c='C1',  marker='|', alpha=0.8)

    zs = arr2.flatten()
    xs = np.zeros_like(zs)
    xs_copy = np.zeros_like(xs)
    ys = np.zeros_like(zs)

    for i in range(arr2.shape[0]):
        _tmp = np.full_like(xs[arr2.shape[1] * i:arr2.shape[1] * (i + 1)], i)
        ys[arr2.shape[1] * i:arr2.shape[1] * (i + 1)] = np.arange(0, arr2.shape[1], 1)
        xs_copy[arr2.shape[1] * i:arr2.shape[1] * (i + 1)] = _tmp

    ax.scatter(xs_copy, ys, zs, c='C2',  marker='|', alpha=0.8)

    zs = arr3.flatten()
    xs = np.zeros_like(zs)
    xs_copy = np.zeros_like(xs)
    ys = np.zeros_like(zs)

    for i in range(arr3.shape[0]):
        _tmp = np.full_like(xs[arr3.shape[1] * i:arr3.shape[1] * (i + 1)], i)
        ys[arr3.shape[1] * i:arr3.shape[1] * (i + 1)] = np.arange(0, arr3.shape[1], 1)
        xs_copy[arr3.shape[1] * i:arr3.shape[1] * (i + 1)] = _tmp

    ax.scatter(xs_copy, ys, zs, c='C3',  marker='|', alpha=0.8)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z(Elevation)')

    fig.suptitle(u'STRM DEM vs Aster GDEM vs Aster L1A Conducted DEM', fontsize=12, color='C0')
    c1_patch = mpatches.Patch(color='C1', label='STRM DEM')
    c2_patch = mpatches.Patch(color='C2', label='Aster GDEM')
    c3_patch = mpatches.Patch(color='C3', label='Aster L1A Conducted DEM')
    plt.legend(handles=[c1_patch, c2_patch, c3_patch])
    plt.show()

    #pickle.dump(fig, file('sinus.pickle', 'w'))
    #fig = pickle.load(open('sinus.pickle', 'rb'))

    #ax.view_init(elev = 89.9, azim = 100.1)
    # gif_filename = 'animated-3d'
    # for n in range(0, 15):
    #     if n >= 10 and n <= 12:
    #         ax.set_xlabel('')
    #         ax.set_ylabel('')  # don't show axis labels while we move around, it looks weird ax.elev = ax.elev-0.5 #start by panning down slowly if n >= 23 and n <= 36: ax.elev = ax.elev-1.0 #pan down faster if n >= 37 and n <= 60: ax.elev = ax.elev-1.5 ax.azim = ax.azim+1.1 #pan down faster and start to rotate if n >= 61 and n <= 64: ax.elev = ax.elev-1.0 ax.azim = ax.azim+1.1 #pan down slower and rotate same speed if n >= 65 and n <= 73: ax.elev = ax.elev-0.5 ax.azim = ax.azim+1.1 #pan down slowly and rotate same speed if n >= 74 and n <= 76:
    #         ax.elev = ax.elev - 0.2
    #         ax.azim = ax.azim + 0.5  # end by panning/rotating slowly to stopping position
    #     if n == 14:  # add axis labels at the end, when the plot isn't moving around
    #         ax.set_xlabel('Weight')
    #         ax.set_ylabel('Height')
    #         ax.set_zlabel('Evaluation')
    #     fig.suptitle(u'Aster GDEM & Aster-scene-drived DEM', fontsize=12, x=0.5, y=0.85)
    #     plt.savefig("{}/{}/img/{}.png".format(FigDir,gif_filename, str(n).zfill(3)), bbox_inches='tight')
    #
    # plt.close()
    # images = [PIL_Image.open(image) for image in glob.glob("{}/{}/*.png".format(FigDir, gif_filename))]
    # file_path_name = "{}/{}.gif".format(FigDir, gif_filename)
    # writeGif(file_path_name, images, duration=0.1)


if __name__ == '__main__':
    #reprojection(STRM_DEM)
    scatter3d_master()
    #resampling(STRM_DEM)